# 导入必要的库
from datasets import load_dataset
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    TrainingArguments,
    Trainer
)
import numpy as np
import evaluate

def tokenize_function(examples):
    """对输入文本进行分词处理
    
    Args:
        examples: 包含文本数据的字典
    Returns:
        tokenized_output: 分词后的结果
    """
    return tokenizer(examples["text"], padding="max_length", truncation=True)

# 计算评估指标
def compute_metrics(eval_pred):
    """计算评估指标
    
    Args:
        eval_pred: 包含预测结果和标签的元组
    Returns:
        metrics: 计算得到的准确率
    """
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    return {"accuracy": (predictions == labels).astype(np.float32).mean().item()}

# 加载Yelp评论数据集
dataset = load_dataset("yelp_review_full")

# 初始化分词器
tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-cased")

# 对数据集进行分词处理
tokenized_datasets = dataset.map(tokenize_function, batched=True)

# 定义训练参数
training_args = TrainingArguments(
    output_dir="test_trainer",
    evaluation_strategy="epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    num_train_epochs=3,
    weight_decay=0.01,
    load_best_model_at_end=True,
    metric_for_best_model="accuracy",
    save_strategy="epoch",
)

# 为了快速测试，选择较小的训练集和评估集
all_train_dataset = tokenized_datasets["train"].shuffle(seed=42).select(range(55000))
train_dataset = all_train_dataset.select(range(50000))
val_dataset = all_train_dataset.select(range(50000, 55000))
test_dataset = tokenized_datasets["test"].shuffle(seed=42).select(range(1000))

# 评估基础模型性能
base_model = AutoModelForSequenceClassification.from_pretrained(
    "google-bert/bert-base-cased",
    num_labels=5,  # Yelp评分为1-5星
    torch_dtype="auto"
)

base_trainer = Trainer(
    model=base_model,
    args=training_args,
    eval_dataset=test_dataset,
    compute_metrics=compute_metrics,
)

print("基础模型在测试集上的性能：")
base_metrics = base_trainer.evaluate()
print(base_metrics)

# 加载预训练模型用于微调
model = AutoModelForSequenceClassification.from_pretrained(
    "google-bert/bert-base-cased",
    num_labels=5,  # Yelp评分为1-5星
    torch_dtype="auto"
)



# 初始化训练器
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=val_dataset,
    compute_metrics=compute_metrics,
)

# 开始训练
trainer.train()

# 在测试集上评估微调后的模型
print("\n微调模型在测试集上的性能：")
fine_tuned_metrics = trainer.evaluate(test_dataset)
print(fine_tuned_metrics)

# 打印性能对比
print("\n模型性能对比：")
print(f"基础模型准确率：{base_metrics['eval_accuracy']:.4f}")
print(f"微调模型准确率：{fine_tuned_metrics['eval_accuracy']:.4f}")
print(f"性能提升：{(fine_tuned_metrics['eval_accuracy'] - base_metrics['eval_accuracy']) * 100:.2f}%")